import wave

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

from codec import Encode
from config import Config


class WaveDataSet(Dataset):
    def __init__(self, ann_path, robust_flag=False):
        super(WaveDataSet, self).__init__()
        self.robust_flag = robust_flag
        self.files = []
        with open(ann_path) as annotation_lines:
            for annotation_line in annotation_lines:
                annotation_line = annotation_line.strip()
                line = annotation_line.split(',')
                self.files.append({"wav": line[0], "target": line[1]})

    def __len__(self):
        return len(self.files)

    def __getitem__(self, index):
        datafiles = self.files[index % len(self.files)]
        wav_path = datafiles["wav"]
        target = datafiles["target"]
        with wave.open(wav_path, 'r') as wf:
            data = wf.readframes(wf.getnframes())
        data, target = Encode(self.robust_flag)(data, [target])
        return data, target


if __name__ == '__main__':
    dst = WaveDataSet(ann_path="../data/annotation.txt")
    loader = DataLoader(dst, batch_size=2, shuffle=False, pin_memory=True)
    for _batch, (_wave, _label) in enumerate(loader):
        print(_batch)
        print(_wave.type())
        print(np.shape(_wave))
        assert (np.shape(_wave) == torch.Size([2, 1, Config["chunk"]*Config["chunks_per_sample"]]))
        print(_label)
